from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.engine.url import URL


class DB:

    TABLE_SQL = ""
    TABLE_COLUMN_SQL = ""

    def __init__(self, **kwargs):
        """
        example: (url='xxx', pool_pre_ping=True, json_serializer=lambda obj: json.dumps(obj, ensure_ascii=False))
        :param kwargs:
        """
        self.engine = None
        self.SessionLocal = None
        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        if not kwargs.get('url'):
            kwargs = {'url': self.init_from_params(**kwargs)}

        self.engine = create_engine(**kwargs)
        self.SessionLocal = sessionmaker(bind=self.engine, autocommit=False, autoflush=True)

    def init_from_params(self, username, password, host, port, database=None, **kwargs):
        url = URL.create(
            drivername=self.DRIVER_NAME,
            username=username,
            password=password,
            host=host,
            port=port,
            database=database,
            **kwargs
        )
        return url

    def get_session(self):
        try:
            s = self.SessionLocal()
            yield s
        finally:
            s.close()

    def verify(self, sql='SELECT 1;'):
        try:
            self.execute_sql(sql)
            return True
        except Exception as e:
            return False

    def execute_sql(self, sql, as_dict=True):
        with self.SessionLocal() as s:
            resp = s.execute(text(sql)).fetchall()
            if as_dict:
               return [dict(row._mapping) for row in resp] 
            return resp

    def get_table(self, database=None):
        """查询数据库下表"""
        resp = self.execute_sql(self.TABLE_SQL.format(database=database))
        return [dict(row._mapping) for row in resp]

    def get_table_column(self, table, database=None):
        """查询表下字段"""
        database = database or self.engine.url.database
        resp = self.execute_sql(self.TABLE_COLUMN_SQL.format(database=database, table=table))
        return [dict(row._mapping) for row in resp]


class MySQL(DB):
    DRIVER_NAME = "mysql+pymysql"
    TABLE_SQL = "SELECT table_name AS 'table' FROM INFORMATION_SCHEMA.TABLES where table_schema='{database}';"
    TABLE_COLUMN_SQL = """
    SELECT COLUMN_NAME AS name, COLUMN_COMMENT AS comment, DATA_TYPE as type FROM INFORMATION_SCHEMA.COLUMNS
    WHERE TABLE_SCHEMA='{database}' AND TABLE_NAME='{table}'
    """

    def __init__(self, **kwargs):
        super(MySQL, self).__init__(**kwargs)


class Oracle(DB):
    DRIVER_NAME = "oracle+cx_oracle"
    TABLE_SQL = """SELECT table_name AS "table" FROM user_tables"""
    TABLE_COLUMN_SQL = """
    select COLUMN_NAME AS "column", DATA_TYPE AS "type" from user_tab_columns where table_name='{table}'
    """

    def __init__(self, **kwargs):
        super(Oracle, self).__init__(**kwargs)

    def verify(self, sql='select 1 from dual'):
        return self.execute_sql(sql)


class Postgres(DB):
    DRIVER_NAME = "postgresql+psycopg2"
    TABLE_SQL = """SELECT tablename AS table FROM pg_tables WHERE schemaname = 'public';"""
    TABLE_COLUMN_SQL = """
        SELECT a.attname as column, format_type(a.atttypid, a.atttypmod) as type, 
        col_description(a.attrelid, a.attnum) as description, a.attnotnull as not_null 
        FROM pg_class as c, pg_attribute as a
        where c.relname = '{table}' and a.attrelid = c.oid and a.attnum>0"""

    def __init__(self, **kwargs):
        super(Postgres, self).__init__(**kwargs)


class MSSQL(DB):
    """microsoft sql server"""

    TABLE_SQL = """"""
    TABLE_COLUMN_SQL = """"""

    def __init__(self, **kwargs):
        super(MSSQL, self).__init__(**kwargs)


class DB2SQL(DB):
    """DB2 SQL"""

    TABLE_SQL = """"""
    TABLE_COLUMN_SQL = """"""

    def __init__(self, **kwargs):
        super(DB2SQL, self).__init__(**kwargs)


db = DB()  # 默认内置连接


if __name__ == '__main__':
    db.init(url='postgresql+psycopg2://aiops:aiops@192.168.101.79:5432/aiops', pool_pre_ping=True)
    postgres = Postgres(url='postgresql+psycopg2://aiops:aiops@192.168.101.79:5432/aiops', pool_pre_ping=True)
    mysql = MySQL(url='mysql+pymysql://root:password@10.1.1.15:3306/mytest')
    oracle = Oracle(url='oracle+cx_oracle://system:oracle@192.168.101.26:1521/?service_name=helowin')

    def test_db():
        s = db.SessionLocal()
        print(s.execute(text('select 1;')).all())
        s.close()

    def test_postgres():
        a = postgres.execute_sql('select 1;')
        b = postgres.get_table('operation_object')
        c = postgres.get_table_column('operation_object')
        print(a, b, c)

    def test_mysql():
        a = mysql.execute_sql('select 1;')
        b = mysql.get_table('mytest')
        c = mysql.get_table_column('test1')
        print(a, b, c)

    def test_oracle():
        a = oracle.execute_sql('select 1 from dual')
        b = oracle.get_table()
        c = oracle.get_table_column('ortest')
        print(a, b, c)


    test_db()
    test_postgres()
    test_mysql()
    test_oracle()
